"""
AI模型配置管理模块
支持本地模型和云端API的配置管理
"""

import os
from typing import Dict, Any, Optional, List
from dataclasses import dataclass
from enum import Enum


class AIModelType(Enum):
    """AI模型类型"""
    OPENAI = "openai"
    CLAUDE = "claude"
    HUGGINGFACE = "huggingface"
    LOCAL_LLM = "local_llm"
    OLLAMA = "ollama"
    BAIDU_WENXIN = "baidu_wenxin"
    ALIBABA_TONGYI = "alibaba_tongyi"
    ZHIPU_GLM = "zhipu_glm"


@dataclass
class AIModelConfig:
    """AI模型配置"""
    model_type: AIModelType
    model_name: str
    api_key: Optional[str] = None
    api_base: Optional[str] = None
    api_version: Optional[str] = None
    max_tokens: int = 2048
    temperature: float = 0.7
    timeout: int = 60
    retry_count: int = 3
    local_model_path: Optional[str] = None
    device: str = "auto"
    
    # 本地模型特有配置
    load_in_8bit: bool = False
    load_in_4bit: bool = False
    torch_dtype: str = "auto"
    trust_remote_code: bool = True


class AIConfigManager:
    """AI配置管理器"""
    
    def __init__(self):
        self.configs: Dict[str, AIModelConfig] = {}
        self.default_configs = self._load_default_configs()
        self.active_model: Optional[str] = None
    
    def _load_default_configs(self) -> Dict[str, Dict[str, Any]]:
        """加载默认配置"""
        return {
            AIModelType.OPENAI.value: {
                "api_base": "https://api.openai.com/v1",
                "model_name": "gpt-3.5-turbo",
                "max_tokens": 2048,
                "temperature": 0.7
            },
            AIModelType.CLAUDE.value: {
                "api_base": "https://api.anthropic.com",
                "model_name": "claude-3-sonnet-20240229",
                "max_tokens": 2048,
                "temperature": 0.7
            },
            AIModelType.HUGGINGFACE.value: {
                "api_base": "https://api-inference.huggingface.co",
                "model_name": "microsoft/DialoGPT-medium",
                "max_tokens": 1024,
                "temperature": 0.7
            },
            AIModelType.LOCAL_LLM.value: {
                "model_name": "chatglm3-6b",
                "max_tokens": 2048,
                "temperature": 0.7,
                "device": "auto"
            },
            AIModelType.OLLAMA.value: {
                "api_base": "http://localhost:11434",
                "model_name": "qwen2:latest",
                "max_tokens": 2048,
                "temperature": 0.7
            },
            AIModelType.BAIDU_WENXIN.value: {
                "api_base": "https://aip.baidubce.com",
                "model_name": "ernie-bot-turbo",
                "max_tokens": 2048,
                "temperature": 0.7
            },
            AIModelType.ZHIPU_GLM.value: {
                "api_base": "https://open.bigmodel.cn/api/paas/v4",
                "model_name": "glm-4",
                "max_tokens": 2048,
                "temperature": 0.7
            }
        }
    
    def add_config(self, name: str, config: AIModelConfig) -> None:
        """添加AI模型配置"""
        self.configs[name] = config
        if not self.active_model:
            self.active_model = name
    
    def get_config(self, name: str) -> Optional[AIModelConfig]:
        """获取AI模型配置"""
        return self.configs.get(name)
    
    def get_active_config(self) -> Optional[AIModelConfig]:
        """获取当前激活的模型配置"""
        if self.active_model:
            return self.get_config(self.active_model)
        return None
    
    def set_active_model(self, name: str) -> bool:
        """设置激活的模型"""
        if name in self.configs:
            self.active_model = name
            return True
        return False
    
    def remove_config(self, name: str) -> bool:
        """删除AI模型配置"""
        if name in self.configs:
            del self.configs[name]
            if self.active_model == name:
                self.active_model = list(self.configs.keys())[0] if self.configs else None
            return True
        return False
    
    def list_configs(self) -> Dict[str, AIModelConfig]:
        """获取所有配置"""
        return self.configs.copy()
    
    def validate_config(self, config: AIModelConfig) -> bool:
        """验证配置有效性"""
        if config.model_type in [AIModelType.OPENAI, AIModelType.CLAUDE, AIModelType.BAIDU_WENXIN]:
            if not config.api_key:
                return False
        
        if config.model_type == AIModelType.LOCAL_LLM:
            if not config.local_model_path and not config.model_name:
                return False
        
        return True


# 全局AI配置管理器实例
ai_config_manager = AIConfigManager()


# 环境变量配置加载
def load_ai_config_from_env() -> None:
    """从环境变量加载AI配置"""
    
    # OpenAI配置
    if os.getenv('OPENAI_API_KEY'):
        openai_config = AIModelConfig(
            model_type=AIModelType.OPENAI,
            model_name=os.getenv('OPENAI_MODEL', 'gpt-3.5-turbo'),
            api_key=os.getenv('OPENAI_API_KEY'),
            api_base=os.getenv('OPENAI_API_BASE', 'https://api.openai.com/v1'),
            max_tokens=int(os.getenv('OPENAI_MAX_TOKENS', 2048)),
            temperature=float(os.getenv('OPENAI_TEMPERATURE', 0.7))
        )
        ai_config_manager.add_config('openai_default', openai_config)
    
    # Claude配置
    if os.getenv('CLAUDE_API_KEY'):
        claude_config = AIModelConfig(
            model_type=AIModelType.CLAUDE,
            model_name=os.getenv('CLAUDE_MODEL', 'claude-3-sonnet-20240229'),
            api_key=os.getenv('CLAUDE_API_KEY'),
            api_base=os.getenv('CLAUDE_API_BASE', 'https://api.anthropic.com'),
            max_tokens=int(os.getenv('CLAUDE_MAX_TOKENS', 2048)),
            temperature=float(os.getenv('CLAUDE_TEMPERATURE', 0.7))
        )
        ai_config_manager.add_config('claude_default', claude_config)
    
    # 本地模型配置
    if os.getenv('LOCAL_MODEL_PATH'):
        local_config = AIModelConfig(
            model_type=AIModelType.LOCAL_LLM,
            model_name=os.getenv('LOCAL_MODEL_NAME', 'chatglm3-6b'),
            local_model_path=os.getenv('LOCAL_MODEL_PATH'),
            device=os.getenv('LOCAL_MODEL_DEVICE', 'auto'),
            max_tokens=int(os.getenv('LOCAL_MODEL_MAX_TOKENS', 2048)),
            temperature=float(os.getenv('LOCAL_MODEL_TEMPERATURE', 0.7)),
            load_in_8bit=os.getenv('LOCAL_MODEL_8BIT', 'false').lower() == 'true',
            load_in_4bit=os.getenv('LOCAL_MODEL_4BIT', 'false').lower() == 'true'
        )
        ai_config_manager.add_config('local_default', local_config)


# 提示词模板配置
class PromptTemplates:
    """提示词模板"""
    
    TABLE_COMMENT_TEMPLATE = """
请根据以下表信息，为该表生成合适的中文注释：

表名: {table_name}
字段信息:
{fields_info}

请分析表的用途，生成一个简洁准确的表注释（不超过50字）。
只返回注释内容，不要其他说明。
"""
    
    FIELD_COMMENT_TEMPLATE = """
请根据以下信息，为字段生成合适的中文注释：

表名: {table_name}
字段名: {field_name}
字段类型: {field_type}
是否为空: {is_nullable}
默认值: {default_value}
其他字段: {other_fields}

请生成一个简洁准确的字段注释（不超过30字）。
只返回注释内容，不要其他说明。
"""
    
    TABLE_RELATIONSHIP_TEMPLATE = """
请根据以下数据库表信息，分析表的用途和与其他表的关系：

当前表: {table_name}
表注释: {table_comment}
字段信息:
{fields_info}

其他相关表:
{related_tables}

请分析：
1. 该表的主要业务用途
2. 与其他表的关系（主表/子表/关联表等）
3. 在业务流程中的位置

请用简洁的语言描述（不超过100字）。
"""
    
    BATCH_ANALYSIS_TEMPLATE = """
请分析以下数据库表结构，识别业务模块和表间关系：

表信息:
{tables_info}

请分析：
1. 主要业务模块
2. 核心业务表
3. 表间依赖关系
4. 数据流向

请用结构化的方式返回分析结果。
"""